{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  OchreLang.FlattenTraits
-- Copyright   :  (c) Stephen Tetley 2014
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Eliminate traits in Top syntax.
--
--------------------------------------------------------------------------------

module OchreLang.FlattenTraits
  ( 

    TraitDefns
  , allTraits
  , flattenTraits


  ) where

import OchreLang.CompilerMon
import OchreLang.SyntaxBase
import OchreLang.TopSyntax


import Control.Applicative
import qualified Data.Map.Strict as Map
import Text.Printf

import Prelude hiding ( lookup )

type TraitName = Ident
type TraitDefns = Map.Map TraitName TraitDefn


type MethodName = Ident
type MethodDefns = Map.Map MethodName MethodDefn


lookup :: MethodName -> MethodDefns -> Maybe MethodDefn
lookup = Map.lookup


methodDecls :: MethodDefns -> [MethodDecl]
methodDecls = Map.foldrWithKey fn []
  where
    fn vid body ac = MethodDecl vid body : ac


--------------------------------------------------------------------------------
-- Trait operations


local :: Trait -> MethodDefns
local (Trait _ (TraitDefn _ ms)) = foldr fn Map.empty ms
  where
    fn (MethodDecl vid body) ac = Map.insert vid body ac


subtraits :: Trait -> Maybe TraitExpr
subtraits (Trait _ (TraitDefn mbE _)) = mbE



exclude :: MethodDefns -> MethodName -> MethodDefns
exclude = flip Map.delete


-- | this is exactly same as Map.union (left biased)
override :: MethodDefns -> MethodDefns -> MethodDefns
override = Map.union

alias :: MethodDefns -> MethodName -> MethodName -> MethodDefns
alias bigM new old = case lookup new bigM of
   Just _ -> Map.delete new bigM        -- annihilate
   Nothing -> case lookup old bigM of
        Just body -> Map.insert new body bigM
        Nothing   -> bigM               -- do nothing, error later...


-- | TODO - surely there is a better way...
--
compose :: MethodDefns -> MethodDefns -> MethodDefns
compose m1 m2 = let left  = Map.difference m1 m2
                    right = Map.difference m2 m1
                in Map.union left right


--------------------------------------------------------------------------------
-- Flattening


allTraits :: Program -> Compiler TraitDefns
allTraits (Program { program_modules = modules }) = 
    return $ foldr fn1 Map.empty $ Map.elems modules
  where
    fn1 (Module { module_traits = traits}) ac = foldr fn2 ac traits
    fn2 (Trait vid defn)                   ac = Map.insert vid defn ac



-- | This is wrong - vis multiple modules, traits might defined
-- in another module...

flattenTraits :: TraitDefns -> Module -> Compiler Module
flattenTraits env modul@(Module { module_instrs = instrs }) =  
    do { new_insts <- mapM (flattenInstr env) instrs
       ; return $ modul { module_instrs = new_insts, module_traits = [] }
       }
                   

flattenInstr :: TraitDefns -> Instr -> Compiler Instr
flattenInstr env inst = 
    do { defns <- flattenMethods env inst
       ; return $ inst { instr_methods   = methodDecls defns
                       , instr_opt_trait = Nothing
                       }
       }



flattenMethods :: TraitDefns -> Instr -> Compiler MethodDefns
flattenMethods env (Instr { instr_opt_trait = trait_exp
                          , instr_methods   = methods} ) = case trait_exp of
    Nothing -> pure $ mdecls 
    Just ta -> override <$> pure mdecls <*> flattenTE env ta
  where
    mdecls = foldr (\(MethodDecl vid body) -> Map.insert vid body) 
                   Map.empty methods


flattenTE :: TraitDefns -> TraitExpr -> Compiler MethodDefns
flattenTE env (NamedTrait s)        = 
    case Map.lookup s env of
      Just t -> flattenT env (Trait s t)
      Nothing -> throwError "FlattenTraits.flattenTE"
                             (printf "Trait lookup failed: %s" (getIdent s))

flattenTE env (TraitComp ta1 ta2)   = 
    compose <$> flattenTE env ta1 <*> flattenTE env ta2
 
flattenTE env (TraitExclude ta vs)  = 
    do { methods <- flattenTE env ta 
       ; return $ foldr (flip exclude) methods vs
       }

flattenTE env (TraitRename ta vs)   =
    do { m1 <-  flattenTE env ta
       ; let fn = \(TraitAlias new old) acc -> alias acc new old
       ; return $ foldr fn m1 vs 
       }


flattenT :: TraitDefns -> Trait -> Compiler MethodDefns
flattenT env t = case subtraits t of
    Just te -> override <$> pure (local t) <*> flattenTE env te
    Nothing -> pure $ local t

